# -*- coding:utf-8 -*-
import os
import logging
from typing import Optional, List, Dict, Any
from neo4j import GraphDatabase, AsyncGraphDatabase
from neo4j.exceptions import ServiceUnavailable, AuthError
from dotenv import load_dotenv

# 确保加载环境变量
load_dotenv()

logger = logging.getLogger(__name__)


class Neo4jConnection:
    """Neo4j数据库连接管理器（单例模式）"""
    
    _instance = None
    _driver = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(Neo4jConnection, cls).__new__(cls)
        return cls._instance
    
    def __init__(self):
        """初始化Neo4j连接"""
        if self._driver is None:
            self.uri = os.getenv('NEO4J_URI', 'bolt://localhost:7687')
            self.username = os.getenv('NEO4J_USERNAME', 'neo4j')
            self.password = os.getenv('NEO4J_PASSWORD', 'password')
            self.database = os.getenv('NEO4J_DATABASE', 'neo4j')
            self._connect()
    
    def _connect(self):
        """建立连接"""
        try:
            self._driver = GraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password),
                max_connection_lifetime=3600,
                max_connection_pool_size=50,
                connection_acquisition_timeout=120
            )
            # 验证连接
            self._driver.verify_connectivity()
            logger.info(f"成功连接到Neo4j数据库: {self.uri}")
        except AuthError as e:
            logger.error(f"Neo4j认证失败: {e}")
            raise
        except ServiceUnavailable as e:
            logger.error(f"Neo4j服务不可用: {e}")
            raise
        except Exception as e:
            logger.error(f"连接Neo4j失败: {e}")
            raise
    
    def close(self):
        """关闭连接"""
        if self._driver:
            self._driver.close()
            logger.info("Neo4j连接已关闭")
    
    def execute_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """
        执行Cypher查询
        
        Args:
            query: Cypher查询语句
            parameters: 查询参数
            
        Returns:
            查询结果列表
        """
        if not self._driver:
            self._connect()
        
        try:
            with self._driver.session(database=self.database) as session:
                result = session.run(query, parameters or {})
                return [record.data() for record in result]
        except Exception as e:
            logger.error(f"执行查询失败: {query}, 错误: {e}")
            raise
    
    def execute_write(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        执行写入操作
        
        Args:
            query: Cypher写入语句
            parameters: 查询参数
            
        Returns:
            执行结果
        """
        if not self._driver:
            self._connect()
        
        try:
            with self._driver.session(database=self.database) as session:
                result = session.run(query, parameters or {})
                summary = result.consume()
                return {
                    "nodes_created": summary.counters.nodes_created,
                    "relationships_created": summary.counters.relationships_created,
                    "properties_set": summary.counters.properties_set
                }
        except Exception as e:
            logger.error(f"执行写入失败: {query}, 错误: {e}")
            raise
    
    def batch_write(self, queries: List[tuple]) -> Dict[str, Any]:
        """
        批量写入数据
        
        Args:
            queries: [(query, parameters), ...] 查询和参数的元组列表
            
        Returns:
            批量执行结果统计
        """
        if not self._driver:
            self._connect()
        
        total_stats = {
            "nodes_created": 0,
            "relationships_created": 0,
            "properties_set": 0,
            "queries_executed": 0
        }
        
        try:
            with self._driver.session(database=self.database) as session:
                for query, parameters in queries:
                    result = session.run(query, parameters or {})
                    summary = result.consume()
                    total_stats["nodes_created"] += summary.counters.nodes_created
                    total_stats["relationships_created"] += summary.counters.relationships_created
                    total_stats["properties_set"] += summary.counters.properties_set
                    total_stats["queries_executed"] += 1
            
            logger.info(f"批量写入完成: {total_stats}")
            return total_stats
        except Exception as e:
            logger.error(f"批量写入失败: {e}")
            raise
    
    def clear_database(self):
        """清空数据库（谨慎使用）"""
        try:
            query = "MATCH (n) DETACH DELETE n"
            self.execute_write(query)
            logger.warning("数据库已清空")
        except Exception as e:
            logger.error(f"清空数据库失败: {e}")
            raise
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取数据库统计信息"""
        try:
            queries = {
                "total_nodes": "MATCH (n) RETURN count(n) as count",
                "total_relationships": "MATCH ()-[r]->() RETURN count(r) as count",
                "node_labels": "CALL db.labels() YIELD label RETURN collect(label) as labels",
                "relationship_types": "CALL db.relationshipTypes() YIELD relationshipType RETURN collect(relationshipType) as types"
            }
            
            stats = {}
            for key, query in queries.items():
                result = self.execute_query(query)
                if result:
                    if key in ["total_nodes", "total_relationships"]:
                        stats[key] = result[0].get("count", 0)
                    else:
                        stats[key] = result[0].get(list(result[0].keys())[0], [])
            
            return stats
        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {}


class AsyncNeo4jConnection:
    """Neo4j异步连接管理器"""
    
    _instance = None
    _driver = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(AsyncNeo4jConnection, cls).__new__(cls)
        return cls._instance
    
    def __init__(self):
        """初始化异步Neo4j连接"""
        if self._driver is None:
            self.uri = os.getenv('NEO4J_URI', 'bolt://localhost:7687')
            self.username = os.getenv('NEO4J_USERNAME', 'neo4j')
            self.password = os.getenv('NEO4J_PASSWORD', 'password')
            self.database = os.getenv('NEO4J_DATABASE', 'neo4j')
    
    async def connect(self):
        """建立异步连接"""
        if self._driver is None:
            try:
                self._driver = AsyncGraphDatabase.driver(
                    self.uri,
                    auth=(self.username, self.password),
                    max_connection_lifetime=3600,
                    max_connection_pool_size=50
                )
                await self._driver.verify_connectivity()
                logger.info(f"成功建立异步连接到Neo4j: {self.uri}")
            except Exception as e:
                logger.error(f"异步连接Neo4j失败: {e}")
                raise
    
    async def close(self):
        """关闭异步连接"""
        if self._driver:
            await self._driver.close()
            logger.info("Neo4j异步连接已关闭")
    
    async def execute_query(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """
        执行异步查询
        
        Args:
            query: Cypher查询语句
            parameters: 查询参数
            
        Returns:
            查询结果列表
        """
        if not self._driver:
            await self.connect()
        
        try:
            async with self._driver.session(database=self.database) as session:
                result = await session.run(query, parameters or {})
                records = await result.data()
                return records
        except Exception as e:
            logger.error(f"异步查询失败: {query}, 错误: {e}")
            raise
    
    async def execute_write(self, query: str, parameters: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        执行异步写入操作
        
        Args:
            query: Cypher写入语句
            parameters: 查询参数
            
        Returns:
            执行结果
        """
        if not self._driver:
            await self.connect()
        
        try:
            async with self._driver.session(database=self.database) as session:
                result = await session.run(query, parameters or {})
                summary = await result.consume()
                return {
                    "nodes_created": summary.counters.nodes_created,
                    "relationships_created": summary.counters.relationships_created,
                    "properties_set": summary.counters.properties_set
                }
        except Exception as e:
            logger.error(f"异步写入失败: {query}, 错误: {e}")
            raise


# 全局连接实例
neo4j_conn = Neo4jConnection()
async_neo4j_conn = AsyncNeo4jConnection()

